"""
智能体基类定义
提供智能体的通用接口和基础功能
"""

from abc import ABC, abstractmethod
from typing import Dict, Any, Optional, List
from enum import Enum
import asyncio
import uuid
from datetime import datetime

from utils.logger import logger
from utils.exceptions import AgentCallError


class AgentStatus(Enum):
    """智能体状态"""
    IDLE = "idle"           # 空闲
    RUNNING = "running"     # 运行中
    WAITING = "waiting"     # 等待中
    ERROR = "error"         # 错误状态
    COMPLETED = "completed" # 完成


class AgentType(Enum):
    """智能体类型"""
    TEXT_QA = "text_qa"          # 文本问答
    OCR = "ocr"                  # OCR识别
    CALCULATOR = "calculator"  # 计算器
    ORCHESTRATOR = "orchestrator" # 协调器
    KNOWLEDGE_BASE = "knowledge_base" # 知识库
    FALLBACK = "fallback"        # 降级方案


class AgentMessage:
    """智能体间消息"""
    
    def __init__(self, sender: str, receiver: str, message_type: str, 
                 content: Any, correlation_id: Optional[str] = None):
        self.message_id = str(uuid.uuid4())
        self.sender = sender
        self.receiver = receiver
        self.message_type = message_type
        self.content = content
        self.correlation_id = correlation_id or self.message_id
        self.timestamp = datetime.now()
        
    def to_dict(self) -> Dict[str, Any]:
        return {
            "message_id": self.message_id,
            "sender": self.sender,
            "receiver": self.receiver,
            "message_type": self.message_type,
            "content": self.content,
            "correlation_id": self.correlation_id,
            "timestamp": self.timestamp.isoformat()
        }


class BaseAgent(ABC):
    """智能体基类"""
    
    def __init__(self, agent_id: str, agent_type: AgentType, name: str):
        self.agent_id = agent_id
        self.agent_type = agent_type
        self.name = name
        self.status = AgentStatus.IDLE
        self.current_task_id: Optional[str] = None
        self.created_at = datetime.now()
        self.last_heartbeat = datetime.now()
        self.last_activity = datetime.now()  # 添加最后活动时间
        self.health_status = "healthy"  # 添加健康状态
        
        # 消息队列
        self.message_queue: asyncio.Queue = asyncio.Queue()
        self.running = False
        self.task_lock = asyncio.Lock()
        
        logger.info(f"智能体初始化: {self.name} ({self.agent_type.value})")
    
    @abstractmethod
    async def process_task(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """处理具体任务 - 子类实现"""
        pass
    
    @abstractmethod
    async def validate_input(self, input_data: Dict[str, Any]) -> bool:
        """验证输入数据 - 子类实现"""
        pass
    
    async def send_message(self, receiver: str, message_type: str, content: Any,
                          correlation_id: Optional[str] = None) -> str:
        """发送消息给其他智能体"""
        message = AgentMessage(
            sender=self.agent_id,
            receiver=receiver,
            message_type=message_type,
            content=content,
            correlation_id=correlation_id
        )
        
        # 通过消息总线发送（这里简化处理）
        # 实际实现中应该通过消息总线或事件总线
        logger.info(f"发送消息: {message.message_id} -> {receiver}")
        return message.message_id
    
    async def receive_message(self, timeout: float = 30.0) -> Optional[AgentMessage]:
        """接收消息"""
        try:
            message = await asyncio.wait_for(self.message_queue.get(), timeout=timeout)
            logger.info(f"接收消息: {message.message_id} from {message.sender}")
            return message
        except asyncio.TimeoutError:
            logger.warning(f"接收消息超时: {self.name}")
            return None
    
    async def execute_task(self, task_id: str, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """执行任务"""
        async with self.task_lock:
            if self.status != AgentStatus.IDLE:
                raise AgentCallError(f"智能体忙: {self.status.value}")
            
            self.status = AgentStatus.RUNNING
            self.current_task_id = task_id
            
            try:
                logger.info(f"开始执行任务: {task_id}")
                
                # 验证输入
                if not await self.validate_input(task_data):
                    raise ValueError("输入数据验证失败")
                
                # 处理任务
                result = await self.process_task(task_data)
                
                self.status = AgentStatus.COMPLETED
                logger.info(f"任务完成: {task_id}")
                
                return {
                    "task_id": task_id,
                    "agent_id": self.agent_id,
                    "status": "success",
                    "result": result,
                    "timestamp": datetime.now().isoformat()
                }
                
            except Exception as e:
                self.status = AgentStatus.ERROR
                logger.error(f"任务执行失败: {task_id}, 错误: {str(e)}")
                
                return {
                    "task_id": task_id,
                    "agent_id": self.agent_id,
                    "status": "error",
                    "error": str(e),
                    "timestamp": datetime.now().isoformat()
                }
            
            finally:
                self.current_task_id = None
                self.last_heartbeat = datetime.now()
    
    async def start(self):
        """启动智能体"""
        self.running = True
        logger.info(f"智能体启动: {self.name}")
        
        # 启动消息处理循环
        asyncio.create_task(self._message_loop())
    
    async def stop(self):
        """停止智能体"""
        self.running = False
        logger.info(f"智能体停止: {self.name}")
    
    async def _message_loop(self):
        """消息处理循环"""
        while self.running:
            try:
                message = await self.receive_message(timeout=1.0)
                if message:
                    await self._handle_message(message)
                    
            except Exception as e:
                logger.error(f"消息处理错误: {str(e)}")
                await asyncio.sleep(1)
    
    async def _handle_message(self, message: AgentMessage):
        """处理接收到的消息"""
        logger.info(f"处理消息: {message.message_type} from {message.sender}")
        
        # 根据消息类型处理
        if message.message_type == "task_request":
            await self._handle_task_request(message)
        elif message.message_type == "task_result":
            await self._handle_task_result(message)
        elif message.message_type == "health_check":
            await self._handle_health_check(message)
        else:
            logger.warning(f"未知消息类型: {message.message_type}")
    
    async def _handle_task_request(self, message: AgentMessage):
        """处理任务请求"""
        task_data = message.content
        task_id = message.correlation_id
        
        result = await self.execute_task(task_id, task_data)
        
        # 发送结果回协调器
        await self.send_message(
            receiver=message.sender,
            message_type="task_result",
            content=result,
            correlation_id=task_id
        )
    
    async def _handle_task_result(self, message: AgentMessage):
        """处理任务结果"""
        # 子类可以重写此方法处理其他智能体的结果
        logger.info(f"收到任务结果: {message.correlation_id}")
    
    async def _handle_health_check(self, message: AgentMessage):
        """处理健康检查"""
        response = {
            "agent_id": self.agent_id,
            "status": self.status.value,
            "last_heartbeat": self.last_heartbeat.isoformat(),
            "current_task": self.current_task_id
        }
        
        await self.send_message(
            receiver=message.sender,
            message_type="health_response",
            content=response,
            correlation_id=message.correlation_id
        )
    
    async def health_check(self) -> Dict[str, Any]:
        """健康检查"""
        return {
            "status": self.health_status,
            "agent_id": self.agent_id,
            "agent_type": self.agent_type.value,
            "last_activity": self.last_activity.isoformat(),
            "timestamp": datetime.now().isoformat()
        }
    
    def get_info(self) -> Dict[str, Any]:
        """获取智能体信息"""
        return {
            "agent_id": self.agent_id,
            "name": self.name,
            "type": self.agent_type.value,
            "status": self.status.value,
            "current_task": self.current_task_id,
            "created_at": self.created_at.isoformat(),
            "last_heartbeat": self.last_heartbeat.isoformat()
        }